/*
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * openFuyao is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package utils

import (
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"

	colocationv1 "openfuyao.com/colocation-service/pkg/apis/colocation/v1"
)

func ToDesiredColocation(podSpec *corev1.PodSpec) bool {
	changed := false
	// 添加混部污点容忍度
	findToleration := false
	for _, toleration := range podSpec.Tolerations {
		if toleration.Key == colocationv1.ColocationNodeLabel && toleration.Operator == corev1.TolerationOpExists && toleration.Effect == corev1.TaintEffectNoExecute {
			findToleration = true
			break
		}
	}

	if !findToleration {
		podSpec.Tolerations = append(podSpec.Tolerations, corev1.Toleration{
			Key:      colocationv1.ColocationNodeLabel,
			Operator: corev1.TolerationOpExists,
			Effect:   corev1.TaintEffectNoExecute,
		})
		klog.Info("add toleration")
		changed = true
	}

	// 添加节点亲和性
	colocationNodeSelectorRequirement := corev1.NodeSelectorRequirement{
		Key:      colocationv1.ColocationTaintLabel,
		Operator: corev1.NodeSelectorOpExists,
	}
	nodeSelectorTerm := corev1.NodeSelectorTerm{
		MatchExpressions: []corev1.NodeSelectorRequirement{
			colocationNodeSelectorRequirement,
		},
	}

	nodeSelector := &corev1.NodeSelector{
		NodeSelectorTerms: []corev1.NodeSelectorTerm{
			nodeSelectorTerm,
		},
	}

	if podSpec.Affinity == nil {
		changed = true
		podSpec.Affinity = &corev1.Affinity{
			NodeAffinity: &corev1.NodeAffinity{
				RequiredDuringSchedulingIgnoredDuringExecution: nodeSelector,
			},
		}
		klog.Info("add node affinity")
		return changed
	}

	if podSpec.Affinity.NodeAffinity == nil {
		changed = true
		podSpec.Affinity.NodeAffinity = &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: nodeSelector,
		}
		klog.Info("add node affinity")
		return changed
	}

	if podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
		changed = true
		podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector
		klog.Info("add node affinity")
		return changed
	}

	findNodeSelectorRequirement := false
	nodeSelectorTerms := []corev1.NodeSelectorTerm{}
	emptyNst := false
	for _, nst := range podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms {
		if len(nst.MatchFields) == 0 && len(nst.MatchExpressions) == 0 {
			klog.Info("find empty nst.......")
			emptyNst = true
			continue
		}

		for _, nsr := range nst.MatchExpressions {
			if nsr.Key == colocationNodeSelectorRequirement.Key && nsr.Operator == colocationNodeSelectorRequirement.Operator {
				findNodeSelectorRequirement = true
				break
			}
		}

		for _, nf := range nst.MatchFields {
			if nf.Key == colocationNodeSelectorRequirement.Key && nf.Operator == colocationNodeSelectorRequirement.Operator {
				findNodeSelectorRequirement = true
				break
			}
		}
		nodeSelectorTerms = append(nodeSelectorTerms, nst)
	}

	if !findNodeSelectorRequirement {
		nodeSelectorTerms = append(nodeSelectorTerms, nodeSelectorTerm)
	}

	if !findNodeSelectorRequirement || emptyNst {
		podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = nodeSelectorTerms
		changed = true
		klog.Info("add node affinity")
	}

	return changed
}

func ToDesiredColocationJob(job *batchv1.Job) bool {
	return ToDesiredColocation(&job.Spec.Template.Spec)
}

func ToDesiredColocationCronJob(cronJob *batchv1.CronJob) bool {
	return ToDesiredColocation(&cronJob.Spec.JobTemplate.Spec.Template.Spec)
}

func ToDesiredNonColocation(podSpec *corev1.PodSpec) bool {
	changed := false
	// 删除混部污点容忍度
	findToleration := false
	tolerations := []corev1.Toleration{}
	for _, toleration := range podSpec.Tolerations {
		if toleration.Key != colocationv1.ColocationNodeLabel {
			tolerations = append(tolerations, toleration)
		} else {
			findToleration = true
			break
		}
	}

	if findToleration {
		klog.Info("delete toleration")
		podSpec.Tolerations = tolerations
		changed = true
	}

	// 删除节点亲和性
	if podSpec.Affinity == nil {
		return changed
	}

	if podSpec.Affinity.NodeAffinity == nil {
		return changed
	}

	if podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
		return changed
	}

	nodeSelectorTerms := podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
	nodeSelectorChanged := false
	for i, nodeSelectorTerm := range nodeSelectorTerms {
		matchExpressions := []corev1.NodeSelectorRequirement{}
		findNodeSelectorRequirement := false
		for _, nodeSelectorRequirement := range nodeSelectorTerm.MatchExpressions {
			if nodeSelectorRequirement.Key == colocationv1.ColocationTaintLabel && nodeSelectorRequirement.Operator == corev1.NodeSelectorOpExists {
				findNodeSelectorRequirement = true
			} else {
				matchExpressions = append(matchExpressions, nodeSelectorRequirement)
			}
		}

		if findNodeSelectorRequirement {
			klog.Info("delete node affinity")
			nodeSelectorTerms[i].MatchExpressions = matchExpressions
			changed = true
			nodeSelectorChanged = true
		}
	}

	if nodeSelectorChanged {
		// 删除空的MatchExpressions
		okNodeSelectorTerm := []corev1.NodeSelectorTerm{}
		nodeSelectorChanged = false
		for _, nodeSelectorTerm := range nodeSelectorTerms {
			if len(nodeSelectorTerm.MatchExpressions) == 0 && len(nodeSelectorTerm.MatchFields) == 0 {
				nodeSelectorChanged = true
			} else {
				okNodeSelectorTerm = append(okNodeSelectorTerm, nodeSelectorTerm)
			}
		}
		if nodeSelectorChanged {
			klog.Info("okNodeSelectorTerm")
			if len(okNodeSelectorTerm) == 0 {
				podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nil
				if len(podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution) == 0 {
					podSpec.Affinity.NodeAffinity = nil
				}
			} else {
				podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = okNodeSelectorTerm
			}
		}
	}

	return changed
}

type StringCheckMode string

const (
	// 常规模式
	NormalMode StringCheckMode = "normal"
	// 自定义模式: 标签Key的格式化校验
	LabelKeyMode StringCheckMode = "label_key_mode"
)
